import redis
import random
from config import config


class RedisHandler:
    def __init__(self):
        self.conn = self.__connection()

    def __connection(self):
        redis_config = config["redis_config"]
        conn = redis.StrictRedis(host=redis_config["host"], port=redis_config["port"], db=redis_config["db"],
                                 decode_responses=True, socket_connect_timeout=3)
        return conn

    # 删除redis数据库里的ip
    def remove_ip(self, ip):
        self.conn.zrem("IP", ip)

    # 获取redis数据库里一共有多少ip
    def get_ip_num(self):
        return self.conn.zcard("IP")

    # 获取ip的端口
    def get_port(self, ip):
        port = self.conn.zscore("IP", ip)
        port = str(port).replace(".0", "")
        return port

    # 添加ip和端口到数据库里
    def add_ip(self, ip, port):
        # nx为true：不存在则set操作才执行
        self.conn.zadd("IP", {ip: port}, nx=True)

    # 列出所有的ip
    def get_all_ip(self):
        return self.conn.zrange("IP", 0, -1)

    # 随机获取一个ip
    def get_random_ip(self):
        end_num = self.get_ip_num()
        num = random.randint(0, end_num)
        random_ip = self.conn.zrange("IP", num, num)
        if not random_ip:
            return "", ""
        random_ip = str(random_ip[0]).replace("b", '').replace("'", "")
        port = self.get_port(random_ip)
        return random_ip, port


if __name__ == '__main__':
    redis = RedisHandler()
    print(redis.get_random_ip())
    # print("over")
